import torch
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from pathlib import Path
from typing import Tuple
import os

class ETTHDataset(Dataset):
    def __init__(self, data_path: str, sequence_length: int, target_horizon: int = 1):
        # 读取ETTh数据集
        df = pd.read_csv(data_path)
        
        # 数据预处理
        self.data = df.select_dtypes(include=[np.float64, np.int64]).values
        self.data = (self.data - self.data.mean(axis=0)) / self.data.std(axis=0)
        
        self.sequence_length = sequence_length
        self.target_horizon = target_horizon
        
    def __len__(self):
        return len(self.data) - self.sequence_length - self.target_horizon + 1
    
    def __getitem__(self, idx):
        x = torch.FloatTensor(self.data[idx:idx + self.sequence_length])
        y = torch.FloatTensor(self.data[idx + self.sequence_length:idx + self.sequence_length + self.target_horizon, 0])
        return x, y

class SafetyHelmetDataset(Dataset):
    def __init__(self, data_dir: str, transform=None):
        self.data_dir = Path(data_dir)
        self.transform = transform or transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        # 读取数据集
        self.images = []
        self.labels = []
        
        for class_id, class_name in enumerate(['no_helmet', 'helmet']):
            class_dir = self.data_dir / class_name
            for img_path in class_dir.glob('*.jpg'):
                self.images.append(img_path)
                self.labels.append(class_id)
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path = self.images[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, self.labels[idx]

def prepare_etth_data(
    data_path: str,
    batch_size: int,
    sequence_length: int,
    train_ratio: float = 0.8
) -> Tuple[DataLoader, DataLoader]:
    """准备ETTh数据集的训练集和验证集"""
    dataset = ETTHDataset(data_path, sequence_length)
    train_size = int(len(dataset) * train_ratio)
    val_size = len(dataset) - train_size
    
    train_dataset, val_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size]
    )
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    return train_loader, val_loader

def prepare_helmet_data(
    data_dir: str,
    batch_size: int,
    train_ratio: float = 0.8
) -> Tuple[DataLoader, DataLoader]:
    """准备安全帽数据集的训练集和验证集"""
    dataset = SafetyHelmetDataset(data_dir)
    train_size = int(len(dataset) * train_ratio)
    val_size = len(dataset) - train_size
    
    train_dataset, val_dataset = torch.utils.data.random_split(
        dataset, [train_size, val_size]
    )
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    
    return train_loader, val_loader 